#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
参数更新工具

供Agent调用的参数更新工具，用于在对话过程中更新洪水预报参数
"""

from langchain.tools import BaseTool
from typing import Dict, Any, Optional, Type
from pydantic import BaseModel, Field
import json
from loguru import logger


class ParameterUpdateInput(BaseModel):
    """参数更新工具的输入模型"""
    reservoir_name: Optional[str] = Field(None, description="水库名称")
    forecast_model: Optional[str] = Field(None, description="预报模型")
    start_time: Optional[str] = Field(None, description="开始时间，格式：YYYY-MM-DD HH:MM")
    forecast_duration: Optional[int] = Field(None, description="预报时长（小时）")
    rainfall_forecast_period: Optional[int] = Field(None, description="降雨预见期（小时）")
    target_variable: Optional[str] = Field(None, description="目标变量")


class ParameterUpdateTool(BaseTool):
    """参数更新工具
    
    允许Agent在对话过程中更新洪水预报参数
    """
    
    name: str = "update_parameters"
    description: str = """
    更新洪水预报参数。当用户提供了新的参数信息时使用此工具。
    
    参数说明：
    - reservoir_name: 水库名称（如：大伙房水库、三峡水库等）
    - forecast_model: 预报模型（如：新安江模型、LSTM模型等）
    - start_time: 开始时间，格式：YYYY-MM-DD HH:MM
    - forecast_duration: 预报时长，单位小时（如：24、48、72）
    - rainfall_forecast_period: 降雨预见期，单位小时（如：12、24、48）
    - target_variable: 目标变量（如：流量、水位、库容）
    
    只传入用户明确提到的参数，未提及的参数不要传入。
    """
    args_schema: Type[BaseModel] = ParameterUpdateInput
    agent_instance: Optional[Any] = Field(default=None, exclude=True)
    
    def __init__(self, agent_instance=None, **kwargs):
        super().__init__(**kwargs)
        object.__setattr__(self, 'agent_instance', agent_instance)
    
    def _run(self, *args, **kwargs) -> str:
        """执行参数更新，兼容位置参数和关键字参数
        
        Args:
            *args: 可能包含参数字典
            **kwargs: 参数字典，包含要更新的参数和LangChain的回调管理器
        Returns:
            str: 更新结果描述
        """
        try:
            if not self.agent_instance:
                return "错误：未设置代理实例"
            
            # 添加调试日志（使用INFO级别以便在前端显示）
            logger.info(f"参数更新工具接收到的args: {args}")
            logger.info(f"参数更新工具接收到的kwargs: {kwargs}")
            
            # 处理LangChain传入的参数格式
            if args:
                # 如果第一个参数是字符串（JSON格式），尝试解析
                if isinstance(args[0], str):
                    try:
                        # 尝试直接解析JSON
                        parsed_params = json.loads(args[0])
                        if isinstance(parsed_params, dict):
                            kwargs.update(parsed_params)
                            logger.info(f"解析JSON参数成功: {parsed_params}")
                    except json.JSONDecodeError as e:
                        # 如果直接解析失败，尝试提取第一个JSON对象
                        try:
                            # 查找第一个完整的JSON对象
                            text = args[0].strip()
                            if text.startswith('{'):
                                # 找到第一个JSON对象的结束位置
                                brace_count = 0
                                end_pos = 0
                                for i, char in enumerate(text):
                                    if char == '{':
                                        brace_count += 1
                                    elif char == '}':
                                        brace_count -= 1
                                        if brace_count == 0:
                                            end_pos = i + 1
                                            break
                                
                                if end_pos > 0:
                                    json_part = text[:end_pos]
                                    parsed_params = json.loads(json_part)
                                    if isinstance(parsed_params, dict):
                                        kwargs.update(parsed_params)
                                        logger.info(f"提取JSON参数成功: {parsed_params}")
                                else:
                                    logger.error(f"无法找到完整的JSON对象: {text[:100]}...")
                            else:
                                logger.error(f"参数不是JSON格式: {text[:100]}...")
                        except Exception as extract_error:
                            logger.error(f"JSON参数提取失败: {extract_error}")
                # 如果第一个参数是字典，直接合并
                elif isinstance(args[0], dict):
                    kwargs.update(args[0])
                    logger.info(f"合并字典参数: {args[0]}")
            
            logger.info(f"处理后的kwargs: {kwargs}")
            
            # 过滤掉LangChain内部参数和None值
            updates = {k: v for k, v in kwargs.items() 
                      if k not in ['run_manager'] and v is not None}
            
            logger.info(f"过滤后的updates: {updates}")
            
            if not updates:
                return "没有需要更新的参数"
            updated_params = []
            for key, value in updates.items():
                if key in ['reservoir_name', 'forecast_model', 'start_time', 
                          'forecast_duration', 'rainfall_forecast_period', 'target_variable']:
                    self.agent_instance.collected_parameters[key] = value
                    updated_params.append(f"{key}: {value}")
                    logger.info(f"更新参数 {key} = {value}")
            if updated_params:
                result = f"已更新参数：{', '.join(updated_params)}"
                # 不在工具返回中判断参数完整性，让Agent基于系统提示中的状态信息来决定下一步
                return result
            else:
                return "没有有效的参数更新"
        except Exception as e:
            logger.error(f"参数更新失败: {str(e)}")
            return f"参数更新失败: {str(e)}"
    
    async def _arun(self, **kwargs) -> str:
        """异步执行参数更新"""
        return self._run(**kwargs)


# if __name__ == "__main__":
#     # 测试用例
#     print("=== 参数更新工具测试 ===")
    
#     # 创建工具实例（不绑定代理）
#     tool = ParameterUpdateTool()
    
#     # 测试参数更新
#     test_params = {
#         "reservoir_name": "大伙房水库",
#         "forecast_model": "新安江模型",
#         "forecast_duration": 48
#     }
    
#     print(f"测试参数: {test_params}")
#     result = tool._run(**test_params)
#     print(f"更新结果: {result}")